import torch
import cv2
from torch.autograd import Variable
from torchvision import transforms,models
from utils.toonnx import to_onnx
import torch.nn as nn
import os
import struct
import torch.nn.functional as F
from PIL import Image
from model.resnet import resnet18
import math
from model.cnn import Net
def resnet_cifar(net, input_data):
    x = net.conv1(input_data)
    x = net.bn1(x)
    x = F.relu(x)
    x = net.maxpool(x)
    x = net.layer1(x)
    x = net.layer2(x)
    x = net.layer3(x)
    x = net.layer4(x)
    x = net.avgpool(x)
    x=x.view(x.size(0),-1)
    return x



def my_forward(model, x):
    mo = nn.Sequential(*list(model.children())[:-1])
    feature = mo(x)
    feature = feature.view(x.size(0), -1)
    #output= model.fc(feature)
    return feature # , output

def pth_test():
    use_cuda = True

    model = models.resnet18(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 10)
    # root = r'D:\proj\button\new_button\best_model_by_138.pth'
    root = './snapshot/params_1.pth'
    xxx = torch.load(root, map_location='cpu')
    model.load_state_dict({k.replace("module.", ""):v for k, v in xxx.items()})
    print('load model')

    print(model)
    model.eval()
    if use_cuda and torch.cuda.is_available():
        model.cuda()


    img = cv2.imread('4_00440.jpg')
    # img = Image.open('test.png').convert('RGB')
    img_tensor = transforms.ToTensor()(img)
    #img_tensor = transforms.Normalize(mean=[],std=[])
    img_tensor = img_tensor.unsqueeze(0)


    if use_cuda and torch.cuda.is_available():
        #prediction1 = resnet_cifar(model, Variable(img_tensor.cuda()))
        prediction2 = model(Variable(img_tensor.cuda()))
        #prediction3 = my_forward(model, Variable(img_tensor.cuda()))
    else:
        #prediction1 = resnet_cifar(model, Variable(img_tensor))
        prediction2 = model(Variable(img_tensor))
        #prediction3 = my_forward(model, Variable(img_tensor))
    # _p = []
    # t_log = 0
    # for i in range(len(x[0])):
    #     t_log += (math.exp(x[0][i]))
    # for i in range(len(x[0])):
    #     p = (math.exp(x[0][i])/t_log)
    #     _p.append(p)
    # xx = sum(_p)
    pred = torch.max(prediction2, 1)[1]
    # t = max(_p)
    print(pred)
    cv2.imshow("image", img)
    cv2.waitKey(0)

def read_bin():
    file_path = 'feature.bin'
    binfile = open(file_path,'rb')
    size = os.path.getsize(file_path)
    pre_data = []
    for i in range(int(size/4)):
        data = binfile.read(4)
        num = struct.unpack('f', data)
        num = list(num)
        pre_data.append('{:.4f}'.format((num[0])))
    binfile.close()
    for i in pre_data:
        print(i)
    num_pre_data = len(pre_data)
    print(num_pre_data)
    return pre_data, num_pre_data

if __name__ == '__main__':
    #read_bin()
    pth_test()
